import os
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from PIL import Image
import requests
from io import BytesIO
import json
from datetime import datetime
from tabulate import tabulate
from colorama import Fore, Style, init

import torch
import torch.nn as nn
from torchvision import transforms

from model import get_model
from train import get_classes, get_data_loaders

# 初始化colorama
init(autoreset=True)

def apply_test_transforms(image):
    """将输入图像转换为模型输入格式"""
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])
    return transform(image)

def predict(model, filepath, classes, device, show_img=False, url=False):
    """使用模型预测图像类别"""
    # 加载图像
    if url:
        response = requests.get(filepath)
        image = Image.open(BytesIO(response.content))
    else:
        image = Image.open(filepath)
    
    # 显示图像（如果需要）
    if show_img:
        plt.figure(figsize=(5, 5))
        plt.imshow(image)
        plt.axis('off')
        plt.show()
    
    # 预处理图像
    image_tensor = apply_test_transforms(image)
    image_batch = torch.unsqueeze(image_tensor, 0).to(device)
    
    # 预测
    model.eval()
    with torch.no_grad():
        output = model(image_batch)
        probabilities = torch.nn.functional.softmax(output, dim=1)
        top5_prob, top5_indices = torch.topk(probabilities, 5)
        
        # 获取预测结果
        top5_classes = []
        for i in range(5):
            idx = top5_indices[0][i].item()
            prob = top5_prob[0][i].item() * 100
            top5_classes.append((classes[idx], prob))
    
    return top5_classes

def format_class_name(class_name):
    """格式化类名，移除前缀和替换连字符"""
    if '.' in class_name and class_name[0].isdigit():
        # 移除类似 "001." 的前缀
        parts = class_name.split('.')
        if len(parts) > 1:
            class_name = '.'.join(parts[1:])
    return class_name.replace('-', ' ').replace('_', ' ')

def test_model(model, test_loader, criterion, classes, device):
    """在测试集上评估模型"""
    test_loss = 0.0
    class_correct = list(0. for i in range(len(classes)))
    class_total = list(0. for i in range(len(classes)))
    
    model.eval()
    
    print("\n" + "="*100)
    print(f"{Fore.CYAN}{Style.BRIGHT}🔍 鸟类分类模型测试 🔍{Style.RESET_ALL}".center(110))
    print(f"{Fore.YELLOW}时间: {Style.BRIGHT}{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}{Style.RESET_ALL}".center(110))
    print("="*100 + "\n")
    
    print(f"{Fore.GREEN}在测试集上评估模型...{Style.RESET_ALL}")
    
    # 使用tqdm显示进度
    progress_bar = tqdm(
        test_loader, 
        desc=f"{Fore.BLUE}{Style.BRIGHT}测试{Style.RESET_ALL}",
        bar_format="{desc}: {percentage:3.0f}%|{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
        ncols=120
    )
    
    # 用于存储所有预测结果的列表
    all_preds = []
    all_targets = []
    
    for data, target in progress_bar:
        data, target = data.to(device), target.to(device)
        
        with torch.no_grad():
            output = model(data)
            loss = criterion(output, target)
            
        test_loss += loss.item() * data.size(0)
        _, pred = torch.max(output, 1)
        
        # 存储预测和目标
        all_preds.extend(pred.cpu().numpy())
        all_targets.extend(target.cpu().numpy())
        
        correct = pred.eq(target.view_as(pred))
        correct_np = correct.cpu().numpy()
        
        # 为每个类别收集统计信息
        for i in range(len(target)):
            label = target[i].item()
            class_correct[label] += correct_np[i].item()
            class_total[label] += 1
    
    # 计算总体测试损失
    test_data_size = sum(class_total)
    test_loss = test_loss / test_data_size
    print(f'\n{Fore.YELLOW}测试损失: {test_loss:.6f}{Style.RESET_ALL}\n')
    
    # 计算总体准确率
    total_correct = sum(class_correct)
    total_samples = sum(class_total)
    overall_acc = 100 * total_correct / total_samples
    
    # 计算Top-5准确率
    top5_correct = 0
    for i in range(len(all_targets)):
        target = all_targets[i]
        output = model(torch.tensor(data[i % len(data)].unsqueeze(0)).to(device))
        _, top5_indices = torch.topk(output, 5)
        if target in top5_indices.cpu().numpy()[0]:
            top5_correct += 1
    top5_acc = 100 * top5_correct / len(all_targets)
    
    # 打印总体性能指标
    print(f"{Fore.GREEN}{Style.BRIGHT}📊 测试结果摘要 📊{Style.RESET_ALL}\n")
    
    performance_table = [
        ["总体准确率", f"{Fore.CYAN}{overall_acc:.2f}%{Style.RESET_ALL} ({int(total_correct)}/{int(total_samples)})"],
        ["Top-5准确率", f"{Fore.CYAN}{top5_acc:.2f}%{Style.RESET_ALL}"],
        ["测试损失", f"{Fore.YELLOW}{test_loss:.4f}{Style.RESET_ALL}"],
        ["测试样本数", f"{total_samples}"],
        ["类别数量", f"{len(classes)}"]
    ]
    
    print(tabulate(performance_table, headers=["指标", "数值"], tablefmt="fancy_grid"))
    
    # 计算每个类别的准确率并排序
    class_accuracies = []
    for i in range(len(classes)):
        if class_total[i] > 0:
            class_acc = 100 * class_correct[i] / class_total[i]
            class_accuracies.append((i, class_acc, class_total[i], class_correct[i]))
    
    # 按准确率排序
    class_accuracies.sort(key=lambda x: x[1], reverse=True)
    
    # 打印每个类别的准确率
    print(f"\n{Fore.MAGENTA}{Style.BRIGHT}🔍 各类别准确率 (前10名) 🔍{Style.RESET_ALL}\n")
    class_table = []
    for i, (idx, acc, total, correct) in enumerate(class_accuracies[:10]):
        class_name = format_class_name(classes[idx])
        class_table.append([
            i+1,
            class_name,
            f"{Fore.CYAN}{acc:.2f}%{Style.RESET_ALL}",
            f"{int(correct)}/{int(total)}"
        ])
    
    print(tabulate(class_table, headers=["排名", "类别", "准确率", "正确/总数"], tablefmt="fancy_grid"))
    
    # 打印表现最差的类别
    print(f"\n{Fore.RED}{Style.BRIGHT}⚠ 各类别准确率 (后10名) ⚠{Style.RESET_ALL}\n")
    class_table = []
    for i, (idx, acc, total, correct) in enumerate(class_accuracies[-10:]):
        class_name = format_class_name(classes[idx])
        class_table.append([
            len(class_accuracies)-9+i,
            class_name,
            f"{Fore.RED}{acc:.2f}%{Style.RESET_ALL}",
            f"{int(correct)}/{int(total)}"
        ])
    
    print(tabulate(class_table, headers=["排名", "类别", "准确率", "正确/总数"], tablefmt="fancy_grid"))
    
    # 保存测试结果到文件
    results_dir = os.path.join(os.getcwd(), 'results')
    os.makedirs(results_dir, exist_ok=True)
    
    # 保存详细结果到JSON
    result_data = {
        "timestamp": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        "overall_accuracy": float(overall_acc),
        "top5_accuracy": float(top5_acc),
        "test_loss": float(test_loss),
        "total_samples": int(total_samples),
        "total_correct": int(total_correct),
        "class_accuracies": [
            {
                "class_name": format_class_name(classes[idx]),
                "accuracy": float(acc),
                "total_samples": int(total),
                "correct_samples": int(correct)
            } for idx, acc, total, correct in class_accuracies
        ]
    }
    
    result_file = os.path.join(results_dir, f'test_results_{datetime.now().strftime("%Y%m%d_%H%M%S")}.json')
    with open(result_file, 'w') as f:
        json.dump(result_data, f, indent=4)
    
    print(f"\n{Fore.GREEN}✓ 测试结果已保存到: {result_file}{Style.RESET_ALL}")
    
    return overall_acc, top5_acc, result_data

def visualize_predictions(model, test_loader, classes, device, num_images=10):
    """可视化模型预测"""
    model.eval()
    
    print(f"\n{Fore.CYAN}{Style.BRIGHT}🖼️ 生成预测可视化 🖼️{Style.RESET_ALL}")
    
    # 获取一批测试图像
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    
    # 将图像移至设备
    images = images.to(device)
    
    # 获取预测
    with torch.no_grad():
        outputs = model(images)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)
        _, preds = torch.max(outputs, 1)
    
    # 将预测和真实标签转回CPU
    preds = preds.cpu().numpy()
    labels = labels.numpy()
    
    # 绘制图像及其预测和真实标签
    fig = plt.figure(figsize=(15, 10))
    for idx in range(min(num_images, len(images))):
        ax = fig.add_subplot(2, 5, idx+1, xticks=[], yticks=[])
        img = images[idx].cpu().numpy().transpose((1, 2, 0))
        
        # 反归一化
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        img = std * img + mean
        img = np.clip(img, 0, 1)
        
        ax.imshow(img)
        
        # 获取预测概率
        prob = probabilities[idx][preds[idx]].item() * 100
        
        # 设置标题
        pred_class = format_class_name(classes[preds[idx]])
        true_class = format_class_name(classes[labels[idx]])
        title = f"预测: {pred_class}\n置信度: {prob:.1f}%\n真实: {true_class}"
        
        # 预测正确显示绿色，错误显示红色
        color = 'green' if preds[idx] == labels[idx] else 'red'
        ax.set_title(title, color=color)
    
    plt.tight_layout()
    
    # 保存可视化结果
    results_dir = os.path.join(os.getcwd(), 'results')
    os.makedirs(results_dir, exist_ok=True)
    viz_file = os.path.join(results_dir, f'prediction_viz_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png')
    plt.savefig(viz_file, dpi=300, bbox_inches='tight')
    
    print(f"{Fore.GREEN}✓ 预测可视化已保存到: {viz_file}{Style.RESET_ALL}")
    
    return viz_file

def main():
    # 设置数据集路径
    dataset_path = "CUB_200_2011/images/"
    
    # 检查数据集是否存在
    if not os.path.exists(dataset_path):
        print(f"{Fore.RED}错误: 数据集路径 {dataset_path} 不存在!{Style.RESET_ALL}")
        return
    
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"{Fore.YELLOW}使用设备: {device}{Style.RESET_ALL}")
    
    # 获取类别
    classes = get_classes(dataset_path)
    num_classes = len(classes)
    print(f"{Fore.YELLOW}类别数量: {num_classes}{Style.RESET_ALL}")
    
    # 获取验证和测试数据加载器
    _, _, test_loader, _ = get_data_loaders(dataset_path, 128, train=False)
    
    # 加载保存的模型
    model_path = 'model/best_model.pth'
    if not os.path.exists(model_path):
        print(f"{Fore.RED}错误: 模型文件 {model_path} 不存在! 请先运行训练脚本。{Style.RESET_ALL}")
        return
    
    # 创建模型
    model = get_model(num_classes)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model = model.to(device)
    
    # 打印模型信息
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"\n{Fore.CYAN}{Style.BRIGHT}📋 模型信息 📋{Style.RESET_ALL}")
    model_info = [
        ["模型类型", "ResNet101"],
        ["总参数量", f"{total_params:,}"],
        ["可训练参数", f"{trainable_params:,} ({trainable_params/total_params*100:.2f}%)"],
        ["非训练参数", f"{total_params - trainable_params:,} ({(total_params-trainable_params)/total_params*100:.2f}%)"],
        ["模型路径", model_path]
    ]
    print(tabulate(model_info, headers=["属性", "值"], tablefmt="fancy_grid"))
    
    # 定义损失函数
    criterion = nn.CrossEntropyLoss()
    
    # 测试模型
    overall_acc, top5_acc, _ = test_model(model, test_loader, criterion, classes, device)
    
    # 可视化一些预测结果
    viz_file = visualize_predictions(model, test_loader, classes, device)
    
    # 测试单个图像
    print(f"\n{Fore.CYAN}{Style.BRIGHT}🔍 单个图像预测示例 🔍{Style.RESET_ALL}")
    
    # 创建示例图像目录
    example_dir = os.path.join(os.getcwd(), 'example_images')
    os.makedirs(example_dir, exist_ok=True)
    
    # 检查是否有示例图像
    example_images = [f for f in os.listdir(example_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]
    
    if example_images:
        sample_image_path = os.path.join(example_dir, example_images[0])
        print(f"使用示例图像: {sample_image_path}")
        
        top5_classes = predict(model, sample_image_path, classes, device, show_img=True)
        
        print(f"\n{Fore.GREEN}{Style.BRIGHT}预测结果:{Style.RESET_ALL}")
        results_table = []
        for i, (class_name, prob) in enumerate(top5_classes):
            formatted_name = format_class_name(class_name)
            results_table.append([i+1, formatted_name, f"{prob:.2f}%"])
        
        print(tabulate(results_table, headers=["排名", "类别", "置信度"], tablefmt="fancy_grid"))
    else:
        print(f"{Fore.YELLOW}提示: 在 'example_images' 目录中没有找到示例图像。{Style.RESET_ALL}")
        print(f"{Fore.YELLOW}您可以将一些鸟类图像放入该目录，然后重新运行测试脚本。{Style.RESET_ALL}")
    
    # 打印测试总结
    print(f"\n{Fore.GREEN}{Style.BRIGHT}✅ 测试完成! ✅{Style.RESET_ALL}")
    print(f"{Fore.CYAN}总体准确率: {overall_acc:.2f}%{Style.RESET_ALL}")
    print(f"{Fore.CYAN}Top-5准确率: {top5_acc:.2f}%{Style.RESET_ALL}")
    print(f"{Fore.CYAN}预测可视化已保存到: {viz_file}{Style.RESET_ALL}")

if __name__ == "__main__":
    main()